# Simple Merger
from asyncio import sleep
from collections import defaultdict
from time import time_ns

from scapy.compat import raw
from scapy.layers.inet import IP, TCP


class Merger:
    def __init__(self, num_inputs, num_outputs):
        # multiple inputs, multiple outputs,
        self.num_inputs = num_inputs
        self.num_outputs = num_outputs
        self.receiver = defaultdict(list)

    def match(self, pkt: bytes, pkt_identifier: int, prev_networks_function_identifier: int):
        content = (pkt, prev_networks_function_identifier)
        if content not in self.receiver[pkt_identifier]:
            self.receiver[pkt_identifier].append(content)
        if len(self.receiver[pkt_identifier]) == self.num_inputs:
            # TODO: special merger function
            data = self.receiver[pkt_identifier]
            del self.receiver[pkt_identifier]
            return pkt, data  # return original pkt
        return False, content


if __name__ == '__main__':
    pkt = b'E\x00\x00-\x00\x01\x00\x00@\x06\x11\xf4\xf7\x89\xe9\xcc$\x91b\xef\x8duM\xf1\x00\x00\x00\x00\x00\x00\x00' \
          b'\x00P\x02 \x00!\xd9\x00\x00abcde '
    count = 0
    merger = Merger(2, 1)
    pkt = IP(src='247.137.233.204', dst='36.145.98.239') / TCP(sport=26213, dport=19953)
    pkt[TCP].add_payload(b'a' * 24)
    pkt = raw(pkt)
    start = time_ns()
    for i in range(100000):
        merger.match(pkt, 1, i % 2)
    print((time_ns() - start) / 100000 / 1000)  # 平均每条指令执行时间
